#!/usr/bin/env julia

"""
    hybrid_ensemble.jl

Hybrid ensemble-variational data assimilation example for GSICoreAnalysis.jl.
This example demonstrates the hybrid method that combines static background 
error covariance (from climatology) with flow-dependent ensemble covariance.

The hybrid background error covariance is:
    B = (1-α) * B_static + α * B_ensemble

where α is the hybrid coefficient (0 = pure 3D-Var, 1 = pure ensemble).

This example demonstrates:
- Setting up ensemble perturbations
- Configuring hybrid ensemble-variational analysis
- Using flow-dependent background error covariance
- Comparing results with pure 3D-Var
- Ensemble inflation and localization concepts
"""

using Pkg
Pkg.activate(@__DIR__ * "/..")

using GSICoreAnalysis
using LinearAlgebra
using Random
using Printf
using Statistics

# Set reproducible random seed
Random.seed!(42)

"""
    generate_ensemble_perturbations(config::AnalysisConfig, truth_state::StateVector{T}) where T

Generate realistic ensemble perturbations around the truth state.
In practice, these would come from ensemble forecasts.
"""
function generate_ensemble_perturbations(config::AnalysisConfig, truth_state::StateVector{T}) where T
    n_ens = config.ensemble_size
    n_state = length(truth_state.values)
    perturbations = zeros(T, n_state, n_ens)
    
    # Generate correlated perturbations
    nx, ny, nz = config.grid_size
    
    for ens in 1:n_ens
        # Create a perturbed ensemble member
        ens_state = StateVector(config)
        assign_state!(ens_state, truth_state)
        
        # Add spatially correlated perturbations
        # Temperature perturbations: correlated in horizontal, varies with height
        for k in 1:nz
            # Create correlation pattern
            base_pattern = randn(nx, ny)
            # Simple spatial smoothing for correlation
            for smooth_iter in 1:2
                for j in 2:ny-1, i in 2:nx-1
                    base_pattern[i,j] = 0.6 * base_pattern[i,j] + 
                                       0.1 * (base_pattern[i-1,j] + base_pattern[i+1,j] + 
                                             base_pattern[i,j-1] + base_pattern[i,j+1])
                end
            end
            
            # Scale by height-dependent factor
            height_factor = 1.0 + 0.5 * sin(π * (k-1) / (nz-1))
            ens_state.t[:,:,k] .+= 1.5 * height_factor * base_pattern
        end
        
        # Wind perturbations: geostrophically balanced with temperature
        for k in 1:nz
            for j in 2:ny-1, i in 2:nx-1
                # Simplified geostrophic balance
                f_coriolis = 1e-4  # Coriolis parameter
                dx = 50000.0  # Grid spacing (m)
                
                # Geostrophic wind from temperature gradient
                dT_dx = (ens_state.t[i+1,j,k] - ens_state.t[i-1,j,k]) / (2*dx)
                dT_dy = (ens_state.t[i,j+1,k] - ens_state.t[i,j-1,k]) / (2*dx)
                
                # Add geostrophic component plus random perturbations
                ens_state.u[i,j,k] += -9.81 * dT_dy / (f_coriolis * 300.0) + 2.0 * randn()
                ens_state.v[i,j,k] += 9.81 * dT_dx / (f_coriolis * 300.0) + 2.0 * randn()
            end
        end
        
        # Humidity perturbations: correlated with temperature
        for k in 1:nz, j in 1:ny, i in 1:nx
            temp_pert = ens_state.t[i,j,k] - truth_state.t[i,j,k]
            humidity_factor = max(0.1, truth_state.q[i,j,k] / 0.01)  # More variability in moist regions
            ens_state.q[i,j,k] += 0.1 * humidity_factor * temp_pert / 2.0 + 0.001 * randn()
            ens_state.q[i,j,k] = max(0.0, ens_state.q[i,j,k])  # Keep positive
        end
        
        # Surface pressure perturbations: correlated with column temperature
        for j in 1:ny, i in 1:nx
            column_temp_pert = mean(ens_state.t[i,j,:] - truth_state.t[i,j,:])
            ens_state.ps[i,j] += -100.0 * column_temp_pert + 200.0 * randn()  # Hydrostatic adjustment
        end
        
        # Store perturbations (difference from truth)
        perturbations[:, ens] .= ens_state.values - truth_state.values
    end
    
    return perturbations
end

"""
    apply_covariance_localization!(perturbations::Matrix{T}, config::AnalysisConfig) where T

Apply spatial localization to ensemble covariances using Gaspari-Cohn function.
"""
function apply_covariance_localization!(perturbations::Matrix{T}, config::AnalysisConfig) where T
    # Simplified localization - in practice would use more sophisticated methods
    nx, ny, nz = config.grid_size
    localization_radius = 5  # Grid points
    
    # This is a placeholder - real localization is much more complex
    # Would typically apply Hadamard product with localization matrix
    
    return perturbations
end

"""
    ensemble_inflation!(perturbations::Matrix{T}, inflation_factor::T) where T

Apply multiplicative inflation to ensemble perturbations.
"""
function ensemble_inflation!(perturbations::Matrix{T}, inflation_factor::T) where T
    perturbations .*= inflation_factor
    return perturbations
end

function main()
    println("=== GSICoreAnalysis.jl Hybrid Ensemble-Variational Example ===")
    println()
    
    # Configuration for hybrid method
    println("Setting up hybrid ensemble configuration...")
    config = AnalysisConfig{Float64}(
        grid_size = (48, 24, 12),       # Smaller grid for ensemble example
        ensemble_size = 25,             # Ensemble size
        hybrid_coeff = 0.75,           # 75% ensemble, 25% static
        use_hybrid = true,
        max_iterations = 40,
        convergence_tol = 1e-4
    )
    
    println("  Grid size: $(config.grid_size)")
    println("  Ensemble size: $(config.ensemble_size)")
    println("  Hybrid coefficient α: $(config.hybrid_coeff)")
    println("  Background covariance: $(100*(1-config.hybrid_coeff))% static + $(100*config.hybrid_coeff)% ensemble")
    println()
    
    # Create truth state
    println("Creating synthetic truth state...")
    true_state = StateVector(config)
    nx, ny, nz = config.grid_size
    
    # Create a weather pattern with a temperature anomaly
    for k in 1:nz, j in 1:ny, i in 1:nx
        # Background temperature field
        height_factor = (k - 1) / (nz - 1)
        lat_factor = cos(π * (j - 1) / (ny - 1))
        true_state.t[i,j,k] = 285.0 - 40.0 * height_factor + 15.0 * lat_factor
        
        # Add a warm anomaly
        center_i, center_j = nx ÷ 2, ny ÷ 2
        distance = sqrt((i - center_i)^2 + (j - center_j)^2)
        anomaly_strength = 8.0 * exp(-distance^2 / 18.0) * (1.0 - height_factor)
        true_state.t[i,j,k] += anomaly_strength
    end
    
    # Wind field in approximate geostrophic balance
    for k in 1:nz, j in 2:ny-1, i in 2:nx-1
        f_coriolis = 1e-4
        dx = 40000.0  # 40 km grid spacing
        
        dT_dx = (true_state.t[i+1,j,k] - true_state.t[i-1,j,k]) / (2*dx)
        dT_dy = (true_state.t[i,j+1,k] - true_state.t[i,j-1,k]) / (2*dx)
        
        true_state.u[i,j,k] = -9.81 * dT_dy / (f_coriolis * 290.0)
        true_state.v[i,j,k] = 9.81 * dT_dx / (f_coriolis * 290.0)
    end
    
    # Other variables
    fill!(true_state.q, 0.008)
    fill!(true_state.ps, 101000.0)
    fill!(true_state.oz, 3e-6)
    fill!(true_state.cw, 0.0)
    fill!(true_state.stl, 285.0)
    fill!(true_state.sti, 271.0)
    
    println("  Created truth with warm temperature anomaly")
    println("  Temperature range: $(extrema(true_state.t)) K")
    println("  Max wind speeds: u=$(maximum(abs.(true_state.u)):.1f), v=$(maximum(abs.(true_state.v)):.1f) m/s")
    println()
    
    # Generate ensemble perturbations
    println("Generating ensemble perturbations...")
    ensemble_perturbations = generate_ensemble_perturbations(config, true_state)
    
    # Apply inflation and localization
    inflation_factor = 1.1  # 10% inflation
    ensemble_inflation!(ensemble_perturbations, inflation_factor)
    apply_covariance_localization!(ensemble_perturbations, config)
    
    # Compute ensemble statistics
    ens_spread_temp = std(ensemble_perturbations[1:nx*ny*nz*3, :], dims=2)  # Temperature portion
    mean_spread = mean(ens_spread_temp)
    
    println("  Ensemble perturbations generated")
    println("  Inflation factor: $(inflation_factor)")
    println("  Mean ensemble spread (temperature): $(mean_spread:.3f) K")
    println()
    
    # Create background state (ensemble mean + error)
    println("Creating background state...")
    background_state = StateVector(config)
    assign_state!(background_state, true_state)
    
    # Add systematic background error
    background_state.t .+= 1.5 * randn(size(background_state.t))
    background_state.u .+= 2.5 * randn(size(background_state.u))
    background_state.v .+= 2.5 * randn(size(background_state.v))
    background_state.q .+= 0.001 * randn(size(background_state.q))
    background_state.ps .+= 300.0 * randn(size(background_state.ps))
    
    bg_temp_rmse = sqrt(mean((background_state.t - true_state.t).^2))
    println("  Background temperature RMSE: $(bg_temp_rmse:.3f) K")
    println()
    
    # Set up hybrid cost function
    println("Setting up hybrid cost function...")
    hybrid_cost = HybridCostFunction(config)
    
    # Initialize static background covariance (simplified diagonal)
    n_state = length(background_state.values)
    static_variances = Float64[]
    
    # Static error variances (from climatology)
    vars_3d = [("u", 2.5^2), ("v", 2.5^2), ("t", 1.5^2), ("q", (0.001)^2), ("oz", (1e-6)^2), ("cw", (0.0005)^2)]
    vars_2d = [("ps", 300.0^2), ("stl", 1.0^2), ("sti", 1.0^2)]
    
    for (var, variance) in vars_3d
        append!(static_variances, fill(variance, nx * ny * nz))
    end
    for (var, variance) in vars_2d
        append!(static_variances, fill(variance, nx * ny))
    end
    
    B_static = Diagonal(static_variances)
    initialize_background_covariance!(hybrid_cost, B_static)
    
    # Initialize ensemble perturbations
    initialize_ensemble!(hybrid_cost, ensemble_perturbations)
    
    # Set background state
    set_background!(hybrid_cost, background_state.values)
    
    println("  Static covariance: diagonal with $(length(static_variances)) elements")
    println("  Ensemble perturbations: $(size(ensemble_perturbations, 2)) members")
    println("  Hybrid coefficient: $(hybrid_cost.hybrid_coefficient)")
    println()
    
    # Create observations
    println("Creating synthetic observations...")
    
    # Observe temperature at specific points
    obs_locations = [(12, 6, 3), (24, 12, 6), (36, 18, 9), (18, 15, 4)]  # (i, j, k)
    n_obs = length(obs_locations)
    observations = zeros(Float64, n_obs)
    obs_error_var = 0.5^2  # 0.5 K observation error
    
    # Simple observation operator for specific points
    function point_observation_operator(state_values::Vector{T}) where T
        state_vec = StateVector{T}(config)
        state_vec.values .= state_values
        
        obs = zeros(T, n_obs)
        for (idx, (i, j, k)) in enumerate(obs_locations)
            obs[idx] = state_vec.t[i, j, k]  # Observe temperature
        end
        return obs
    end
    
    # Generate observations with noise
    true_obs = point_observation_operator(true_state.values)
    for i in 1:n_obs
        observations[i] = true_obs[i] + sqrt(obs_error_var) * randn()
    end
    
    R = Diagonal(fill(obs_error_var, n_obs))
    set_observations!(hybrid_cost, observations, point_observation_operator)
    hybrid_cost.observation_covariance = R
    
    println("  Observations at $(n_obs) points")
    println("  Observation error std: $(sqrt(obs_error_var)) K")
    println("  Innovation (O-B) mean: $(mean(observations - point_observation_operator(background_state.values)):.3f) K")
    println()
    
    # Set up control vector and solver
    println("Setting up hybrid analysis...")
    control_vector = ControlVector(config)
    fill!(control_vector.values, 0.0)
    
    # Use more iterations for hybrid method
    solver = PCGSolver(config, 
                      max_iterations = 50,
                      tolerance = 1e-5,
                      line_search_iterations = 3)
    
    # Initial cost evaluation
    initial_cost = evaluate_cost(hybrid_cost, background_state.values)
    initial_gradient = compute_gradient(hybrid_cost, background_state.values)
    initial_grad_norm = norm(initial_gradient)
    
    println("  Initial cost: $(initial_cost)")
    println("  Initial gradient norm: $(initial_grad_norm)")
    println()
    
    # Run hybrid analysis
    println("Starting hybrid ensemble-variational analysis...")
    println("=" ^ 60)
    
    result = minimize_cost_function(hybrid_cost, control_vector, solver)
    
    println("=" ^ 60)
    println("Hybrid analysis complete!")
    println()
    
    # Compare with pure 3D-Var
    println("Running comparison with pure 3D-Var...")
    
    # Set up pure 3D-Var
    pure_3dvar_config = AnalysisConfig{Float64}(
        grid_size = config.grid_size,
        ensemble_size = 0,
        hybrid_coeff = 0.0,
        use_hybrid = false,
        max_iterations = config.max_iterations,
        convergence_tol = config.convergence_tol
    )
    
    pure_cost = CostFunction(pure_3dvar_config)
    initialize_background_covariance!(pure_cost, B_static)
    set_background!(pure_cost, background_state.values)
    set_observations!(pure_cost, observations, point_observation_operator)
    pure_cost.observation_covariance = R
    
    pure_control = ControlVector(pure_3dvar_config)
    fill!(pure_control.values, 0.0)
    
    pure_solver = PCGSolver(pure_3dvar_config, 
                           max_iterations = 50,
                           tolerance = 1e-5)
    
    pure_result = minimize_cost_function(pure_cost, pure_control, pure_solver)
    
    println("Pure 3D-Var complete!")
    println()
    
    # Results comparison
    println("Results Comparison:")
    println("=" ^ 50)
    println()
    
    println("Hybrid Ensemble-Variational Analysis:")
    println("  Converged: $(result.converged)")
    println("  Iterations: $(result.iterations)")
    println("  Final cost: $(result.final_cost)")
    println("  Final gradient norm: $(result.final_gradient_norm)")
    println("  Cost reduction: $(initial_cost - result.final_cost) ($(((initial_cost - result.final_cost) / initial_cost * 100):.2f)%)")
    println()
    
    println("Pure 3D-Var Analysis:")
    println("  Converged: $(pure_result.converged)")
    println("  Iterations: $(pure_result.iterations)")
    println("  Final cost: $(pure_result.final_cost)")
    println("  Final gradient norm: $(pure_result.final_gradient_norm)")
    cost_3dvar_initial = evaluate_cost(pure_cost, background_state.values)
    println("  Cost reduction: $(cost_3dvar_initial - pure_result.final_cost) ($(((cost_3dvar_initial - pure_result.final_cost) / cost_3dvar_initial * 100):.2f)%)")
    println()
    
    # Convergence comparison
    println("Convergence History Comparison:")
    println("-" ^ 50)
    println("Iter    Hybrid Cost      3D-Var Cost")
    for i in 1:min(10, min(length(result.cost_history), length(pure_result.cost_history)))
        @printf("%3d    %12.6e   %12.6e\n", i-1, result.cost_history[i], pure_result.cost_history[i])
    end
    println()
    
    # Analysis quality assessment
    println("Analysis Quality Assessment:")
    println("-" ^ 30)
    
    # For this example, we'll compare the innovations at observation points
    hybrid_innovations = observations - point_observation_operator(background_state.values .+ result.solution.values)
    pure_innovations = observations - point_observation_operator(background_state.values .+ pure_result.solution.values)
    background_innovations = observations - point_observation_operator(background_state.values)
    
    println("Innovation Statistics (O - A):")
    println("  Background RMS innovation: $(sqrt(mean(background_innovations.^2)):.3f) K")
    println("  Hybrid analysis RMS innovation: $(sqrt(mean(hybrid_innovations.^2)):.3f) K")
    println("  Pure 3D-Var RMS innovation: $(sqrt(mean(pure_innovations.^2)):.3f) K")
    println()
    
    println("Innovation Reduction:")
    hybrid_reduction = (sqrt(mean(background_innovations.^2)) - sqrt(mean(hybrid_innovations.^2))) / sqrt(mean(background_innovations.^2)) * 100
    pure_reduction = (sqrt(mean(background_innovations.^2)) - sqrt(mean(pure_innovations.^2))) / sqrt(mean(background_innovations.^2)) * 100
    println("  Hybrid method: $(hybrid_reduction:.1f)%")
    println("  Pure 3D-Var: $(pure_reduction:.1f)%")
    println()
    
    # Summary
    println("Summary:")
    println("  The hybrid ensemble-variational method combines the benefits of")
    println("  static climatological covariances with flow-dependent ensemble")
    println("  covariances. This example shows how ensemble information can")
    println("  improve the analysis, especially in regions with strong gradients")
    println("  or unusual weather patterns.")
    println()
    
    if hybrid_reduction > pure_reduction
        println("  ✓ Hybrid method shows superior performance in this case")
    else
        println("  → Pure 3D-Var performed better (possibly due to limited ensemble size)")
    end
    
    println()
    println("Hybrid ensemble-variational example completed!")
    
    return (hybrid_result = result, pure_result = pure_result)
end

# Run the example if called directly
if abspath(PROGRAM_FILE) == @__FILE__
    results = main()
end